import numpy as np
import math

def Rosenbroke(xx):
    """
    rosenbroke函数实现
    param:输入向量
    """
    fsum = 0

    for i in range(len(xx)):
        if i == 0:
            continue
        else:
            fsum = fsum + 100*((xx[i]-xx[i-1]**2)**2)+ (xx[i-1]-1)**2
    return fsum

def dRosenbroke(xx):
    """
    rosenbroke导函数
    param:输入向量
    """
    temp = []
    for i in range(len(xx)):
        if i == 0:
            fx = 400*(xx[i]**3) + 2*(xx[i]-1) - 400*xx[i+1]*xx[i]
            temp.append(fx)
        elif i == len(xx)-1:
            fx = 200*xx[i] - 200*(xx[i-1]**2)
            temp.append(fx)
        else:
            fx = 202*xx[i]-200*(xx[i-1]**2) + 100*(4*(xx[i]**3)-4*(xx[i+1]*xx[i])) -2
            temp.append(fx)
    array = np.array(temp)
    return array

def ddRosenbroke(xx):
    """
    rosenbroke海森阵
    param:输入向量
    """
    Hessian = np.zeros((len(xx),len(xx)))
    for i in range(len(xx)):
        if i == 0:
            Hessian[i,i] = 202 - 400 * xx[i+1]  + 1200 * (xx[i]**2)
            Hessian[i,i+1] = -400 * xx[i]
        elif i == len(xx) -1:
            Hessian[i,i-1] = -400 * xx[i-1]
            Hessian[i,i] = 200
        else:
            Hessian[i,i-1] = -400*xx[i-1]
            Hessian[i,i+1] = -400*xx[i]
            Hessian[i,i] = 202 - 400* xx[i+1]  + 1200 *(xx[i]**2)

    return Hessian

if __name__ == "__main__":
    """
    test
    """
    xx = [1,4,5,3,6,8]
    print(ddRosenbroke(xx))